import json
import time
from rag_chain import RAGChain
from text_processor import TextVectorizer

# 配置信息
TEST_FILE_PATH = 'test_search_info.txt'
TEST_PHONE = '13812345678'
EXPECTED_COMPANY = '上海科技发展有限公司'

# 读取测试文件内容
def read_test_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            print(f'Successfully read {file_path}')
            print(f'File content length: {len(content)} characters')
            return content
    except Exception as e:
        print(f'Error reading {file_path}: {e}')
        return None

# 初始化向量器和RAG链
def init_rag_chain():
    try:
        vectorizer = TextVectorizer()
        rag_chain = RAGChain(vectorizer)
        print('Successfully initialized RAG chain')
        return rag_chain
    except Exception as e:
        print(f'Error initializing RAG chain: {e}')
        return None

# 添加文档到向量器（使用vectorize_and_store方法）
def add_document_to_vectorizer(rag_chain, content, doc_id='test_doc'):
    try:
        # 使用vectorize_and_store方法添加文档
        if hasattr(rag_chain.vectorizer, 'vectorize_and_store'):
            # 查看该方法需要的参数
            print(f'Method signature for vectorize_and_store: {rag_chain.vectorizer.vectorize_and_store.__doc__}')
            # 按照方法要求传递参数：text, metadata, chunk_size等
            metadata = {'doc_id': doc_id}
            # 使用默认的chunk_size和chunk_overlap
            rag_chain.vectorizer.vectorize_and_store(content, metadata)
            print(f'Successfully added document to vectorizer using vectorize_and_store')
            return True
        else:
            print(f'Available methods for TextVectorizer: {dir(rag_chain.vectorizer)}')
            raise AttributeError('TextVectorizer has no vectorize_and_store method')
    except Exception as e:
        print(f'Error adding document to vectorizer: {e}')
        return False

# 测试电话号码检索
def test_phone_search(rag_chain, phone_number):
    try:
        print(f'Searching for phone number: {phone_number}')
        result = rag_chain.run(phone_number)
        print(f'Search completed')
        return result
    except Exception as e:
        print(f'Error during search: {e}')
        return None

# 验证搜索结果
def verify_search_result(result, expected_company):
    if not result:
        return False, 'No search result'

    answer = result.get('answer', '')
    print(f'Search answer: {answer}')

    # 检查是否包含预期的公司名称
    if expected_company in answer:
        return True, f'Success! Found {expected_company} in search result'
    else:
        # 检查相关文档
        documents = result.get('documents', [])
        for doc in documents:
            if expected_company in doc.get('content', ''):
                return True, f'Success! Found {expected_company} in related documents'
        return False, f'Failed to find {expected_company} in search result or related documents'

# 主函数
def main():
    print(f'Starting local test for phone number {TEST_PHONE}')

    # 读取测试文件
    file_content = read_test_file(TEST_FILE_PATH)
    if not file_content:
        print('Exiting test due to file read failure')
        return

    # 初始化RAG链
    rag_chain = init_rag_chain()
    if not rag_chain:
        print('Exiting test due to RAG chain initialization failure')
        return

    # 添加文档到向量器
    add_success = add_document_to_vectorizer(rag_chain, file_content)
    if not add_success:
        print('Exiting test due to document addition failure')
        return

    # 等待索引完成
    print('Waiting for document to be indexed...')
    time.sleep(2)  # 等待2秒让向量器有时间处理文档

    # 测试搜索
    search_result = test_phone_search(rag_chain, TEST_PHONE)
    if not search_result:
        print('Exiting test due to search failure')
        return

    # 验证结果
    success, message = verify_search_result(search_result, EXPECTED_COMPANY)
    print(message)

    # 输出完整搜索结果供调试
    print('\nFull search result:')
    print(json.dumps(search_result, ensure_ascii=False, indent=2))

if __name__ == '__main__':
    main()